//      cds_vom.h
//      
//      Copyright 2012 tqlong <tqlong@espada>
//      
//      This program is free software; you can redistribute it and/or modify
//      it under the terms of the GNU General Public License as published by
//      the Free Software Foundation; either version 2 of the License, or
//      (at your option) any later version.
//      
//      This program is distributed in the hope that it will be useful,
//      but WITHOUT ANY WARRANTY; without even the implied warranty of
//      MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//      GNU General Public License for more details.
//      
//      You should have received a copy of the GNU General Public License
//      along with this program; if not, write to the Free Software
//      Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
//      MA 02110-1301, USA.

#ifndef CODING_SEQUENCE_VARIABLE_ORDER_MODEL_H
#define CODING_SEQUENCE_VARIABLE_ORDER_MODEL_H

#include "helper/stl_helper.h"
#include "helper/dna_processing.h"
#include "helper/ublas_helper.h"
#include "helper/suffix_array.h"

class vom_node {
 public:
  char c_; // character at this node
  const std::string* alphabet_;
  std::vector<vec_t> f_; // frequency table at each frame, empty if pruned
  std::vector<vom_node*> child_; // child for each character, NULL if pruned
  vom_node* parent_;
  vi_t n_occurence_; // count pretext occurence at each frame
  vec_t lambda_; // weight for interpolated models at each frame
  int depth_;
 public:
  vom_node(char c, const std::string* alphabet, vom_node* parent, int periodic);
  ~vom_node();
  void set_frequency(const sastring_vec_t& train_set);
  void set_rand_frequency();
  void normalize();
  std::string str() const;
  int periodic() const;
  bool has_table(int pr) const;
  int alphabet_size() const;
  char alphabet(int j) const;

  void copy(vom_node* N);

  vec_t prob_imm(int frame);
};

class VOM {
 public:
  std::string alphabet;
  vom_node* root_;
  double eta_;
  int T_;
  int periodic_;

  int count_;
  int depth_;
  vi_t count_p_, depth_p_;

  double total_node_;
  double current_progress_;
		
  VOM(const std::string& alpha, int periodic);
  ~VOM();

  void build_fom(const sastring_vec_t& train_set, int order);
  void build_fom_flat(const sastring_vec_t& train_set, int order);

  int pos(char c) const;
  char ch(int pos) const;
		
  double log_prob(const std::string& s) const;
		
  void print_node(std::ostream& os) const;
  void load_node(std::istream& is);
		
  void grow(const sastring_vec_t& train_set, double eta);
  void prune(const sastring_vec_t& train_set, double eta);
  void prune_lambda(const sastring_vec_t& train_set, double eta);

  void simulate(int len, std::string& s);
    
  void depth(); // calculate depth, number of nodes for each frame
  void stats_print(const std::string& statsFile) const;

  // IMM
  void compute_lambda_xi_square();
  void compute_lambda_deleted_interpolation();
  double log_prob_lambda(const std::string& s) const;

  // Copy model
  void copy(VOM& m);
 protected:
  void build_fom_node(vom_node* N, const sastring_vec_t& train_set, int order);
  void accumulate_occurence(const std::string& s, int p, int order);
  void print_node(std::ostream& os, vom_node* N, int depth) const;
  void load_node(std::istream& is, vom_node** N, vom_node* parent, int depth);

  double log_prob(const std::string& s, int pos) const;
  const vec_t& f_table(const std::string& s, int pos) const;
  vom_node* find_node(const std::string& s, int pos) const;
  
  void compute_lambda_xi_square(vom_node* N, const std::vector<vec_t>& imm_f);
		
  void grow_node(vom_node* N, const sastring_vec_t& train_set);
  void try_grow(vom_node* N, int pos, const sastring_vec_t& train_set);
  void prune_node(vom_node** N, const sastring_vec_t& train_set);
  void prune_lambda_node(vom_node** N, const sastring_vec_t& train_set);
    
  void depth_node(vom_node* N, int depth);
};


namespace vom {
  struct node {
    const str_t* alphabet_;         // alphabet ACGT
    int periodic_;                  // number of frames
    char ch_;

    node* parent_;                  // tree structure, parent (up)
    std::vector<node*> children_;   // and children (down)

    std::vector<vec_t> count_; // count c_kA c_kC ... at each frame
    std::vector<vec_t> count1_; // count c_kA c_kC ... at each frame (validation set for deleted interpolation)
    std::vector<vec_t> f_;     // frequency A|c_k C|c_k ... at each frame
    vec_t lambda_;             // interpolating parameter at each frame
    int depth_;                // depth of this node
    bool count2_added_;        // technical variable for building VOM

    node(const str_t* alphabet, int periodic, char ch, node* parent);
    ~node() { BOOST_FOREACH(node* child, children_) delete child; }

    int periodic() const { return periodic_; }
    bool has_table(int frame) const { return sz(f_[frame]) == alphabet_size(); }
    int alphabet_size() const { return sz(*alphabet_); }
    char pos2ch(int j) const { return alphabet_->at(j); }
    int ch2pos(char c) const { 
      for (int j = 0; j < alphabet_size(); j++) if (pos2ch(j) == c) return j; 
      return -1;
    }

    void copy(node* N);
    vec_t prob_imm(int frame);
    str_t str();
  };

  typedef std::vector<node*> node_vec_t;

  class vom {
  public:
    str_t alphabet_;
    int periodic_;

    node* root_;

    int count_, depth_;
    vi_t count_frame_, depth_frame_;

    int total_node_, current_progress_;

  public:
    vom(const str_t& alphabet, int periodic);
    ~vom() { delete root_; }
    
    // build tree
    void build_fom_flat(const Sequences& train, int order);
    void build_vom_flat(const Sequences& train, int Tmin);
    int node_pos(const str_t& s, node_vec_t* lv);

    // compute log likelihood
    double log_likelihood(const str_t& test);
    double log_likelihood_lambda(const str_t& s);

    // compute statistics
    void compute_stats();
    void compute_stats_lambda();
    void print_stats(const str_t& outFile);

    // pruning
    void prune(double eta);
    void prune_lambda(double eta);

    // interpolated markov models (lambda weights)
    void compute_lambda_chi_square(int T);
    void compute_lambda_deleted_interpolation(const Sequences& train, double delta);
    void populate_next_level(node_vec_t* lv, node_vec_t* next_lv);
    void split_to_bin(node_vec_t* lv, std::vector<node_vec_t>& bins, int frame, double delta);
    double optimize_bin(node_vec_t& bin, int frame);
    double bin_func(node_vec_t& bin, int frame, double lambda);

    // copy model
    void copy(vom& m);
  };
}

#endif
